#include "MultibodySystem.h"
#include "Force.h"
#include "FastOps.h"
#include "Joint.h"

#include <assert.h>


void UnableToAllocateMemory(){
        cout<<"\nNot successful in allocating required memory. Please email main() to kishor8dm@gmail.com "<<
        "so that this problem can be fixed. Thank you. \n";
        exit(0);
}

using namespace std;
using namespace RSIM;

////////////////////////////////////////////////////////////////////////
//		MultibodySystem
////////////////////////////////////////////////////////////////////////

void MultibodySystem::setTopologyMaps(){
        if(!isTopologyMapSet_){
                isTopologyMapSet_ = true;
                if(Stage::CurrentStage==Stage::MODEL){
                        
                        // initialize iterator for BodyArray
                        BodyArray_.finalize();
                        
                        // initialize iterator for child array of each node
                        for(int i=0;i<nNodes_;++i){
                                BodyArray_[i]->Child_.finalize();
                        }
                        
                        // create joints 
                        this->initJoints();
                        
                        // allocate space and calculate body fixed transform FM.
                        // this remains fixed for the rest of the simulation. 
                        // betak is initialized to FM. This is a fair assumption
                        // since at this stage, state has not been defined and q's
                        // of the system are presumably zero.
                        this->initTransformFM();
                        
                        // Iniitalize the connectivity Map for this system
                        this->initConnectivityMap();
                        
                        Stage::CurrentStage= Stage::TOPOLOGYMAPS;
                }
                else if(Stage::CurrentStage==Stage::EMPTY){
                        cout<<"\nMultibodySystem::setTopologyMaps()\n";
                        Stage::error();
                }
        }
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::initConnectivityMap(){
        List <int> body_to_root;
        RigidBodyNode *nd;
        for(int i=0;i<nNodes_;++i){
                // clear list.
                body_to_root.clear();
                nd = BodyArray_[i];
                do{	
                        body_to_root.addNode(nd->NodeId_);
                        nd = nd->Parent_;			
                }while(nd);
                body_to_root.finalize();
                cMap_.addNode(body_to_root);
        }
        cMap_.finalize();
        for(int i=0;i<nNodes_;++i)
                cMap_[i].finalize();
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::initJoints(){
        for(int i=0;i<nNodes_;++i){
                BodyArray_[i]->data_->initJoint(i);
        }
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::initTransformFM(){
        
        // Root body does not have parent. 
        BodyArray_[0]->data_->wgetpChildID() = -1;
        {
                // TBD
                double tmp[] = {0.0,0.0,0.0};
                BodyArray_[0]->data_->setBetaK(tmp);
                // TBD                                 
        }
        
        for(int i=0;i<nNodes_;++i){

                const int& len = BodyArray_[i]->Child_.getLength();
                const InvTransform& T_FB = InvTransform(BodyArray_[i]->data_->getTransformBF());
                
                BodyArray_[i]->data_->setBodyRK(T_FB.P().Ptr());

                if(len){
                        BodyArray_[i]->data_->allocateSpaceForTransformFM(len);
                        Transform *FM = BodyArray_[i]->data_->wgetBodyTransformFM();

                        for(int j=0;j<len;++j){
                                
                                const Transform& T_PM = BodyArray_[i]->Child_[j]->data_->getTransformPM();

                                FM[j].setR(T_FB.R()*T_PM.R());
                                FM[j].setP(T_FB.P() + T_FB.R()*T_PM.P());
                                                        
                                BodyArray_[i]->Child_[j]->data_->wgetpChildID()=j;
                                
                                // TBD
                                BodyArray_[i]->Child_[j]->data_->setBetaK(FM[j].P().Ptr());
                                // TBD                                                                
                        }
                }
        }
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::updSystemVelocity(const State& state){
        Root_->data_->updRelVelocityOfRoot(state);
        recursiveVelocityUpdate(state, Root_);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::updSystemAcceleration(State& state){
        this->sweepBackwardsABA(state);
        this->sweepForwardABA(state);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::calcSystemAcceleration(State& state){
        this->updStateFromPrescribedMotion(state);
        this->updSystemPosition(state);
        this->updSystemVelocity(state);
        this->updSystemAcceleration(state);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::recursiveVelocityUpdate(const State& state, RigidBodyNode *Nd){
        int len = Nd->Child_.getLength();
        for(int i=0;i<len;++i){
                RigidBodyNode *Child = Nd->Child_.getNode(i);
                Child->data_->updSpatialVelocity(state, Nd, i);
                this->recursiveVelocityUpdate(state, Child);
        }
}

////////////////////////////////////////////////////////////////////////

RigidBodyNode* MultibodySystem::setRoot(const RigidBody& Data){
        
        // Update stage of the system
        Stage::CurrentStage = Stage::MODEL;
        
        // Allocate space
        Root_ = new RigidBodyNode(nNodes_++,Data);
        
        // Store a pointer to the newly created body
        BodyArray_.addNode(Root_);
        
        return Root_;
}

////////////////////////////////////////////////////////////////////////

RigidBodyNode* MultibodySystem::addNodeTo(RigidBodyNode *ParentNode, const RigidBody& Data){
        if(ParentNode==NULL){
                std::cout<<"MultibodySystem::addNodeTo(RigidBodyNode *, const RigidBody&): ParentNode is a Null Pointer. Exiting...\n";
                exit(0);
        }
        else{
                RigidBodyNode* tmp = new RigidBodyNode(ParentNode,nNodes_++,Data);
                (ParentNode->Child_).addNode(tmp);
                BodyArray_.addNode(tmp);
                return tmp;
        }
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::updJointTransforms(const State& state){
        for(int i=0;i<nNodes_;++i)
                BodyArray_[i]->updJointTransform(state.getQForNode(i));
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::updSystemPosition(const State& state){	
        // update the rotation matrics R_FM for each joint
        // update betak if applicable
        // This won't do anything to betak in future.
        updJointTransforms(state);
        
        // Initialize the R_FN matrix for the first body
        Root_->data_->setRotationNF();
        
        // Recursively update all the rotation matrices for the system
        recursiveUpdSystemRotationMatrixNF(state, Root_);
        
        // Update the nbeta vector for the entire system
        recursiveUpdSystemNBetak(state, Root_);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::recursiveUpdSystemNBetak(const State& state, RigidBodyNode *Nd){
        const RigidBodyNode *parent = Nd->Parent_;
        
        // betak vector which is to be updated.
        Vect3& nbetak = Nd->data_->wgetnbetak();
        
        // Joint connecting this node to the parent.
        const Transform& T_MF = Nd->data_->getJointTransformMF();        
        
        double tmp_vec3[3];
        
        if(parent){
                const int& child_index = Nd->data_->getpChildID();
                const Transform *FM = parent->data_->getBodyTransformFM();
                const Rotation& R_NF = parent->data_->getRotationNF();                                
                
                // Conert joint transform FM to parent F basis
                FastOps::mult_mat33_vec3(FM[child_index].R().Ptr(),
                                        T_MF.P().Ptr(),
                                        tmp_vec3
                                        );
                // Add vector body FM to the joint transform
                const double *ptr = FM[child_index].P().Ptr();
                tmp_vec3[0] += ptr[0]; tmp_vec3[1] += ptr[1];tmp_vec3[2] += ptr[2];
                
                // Convert teh result form F basis of the body to the newtonian frame
                FastOps::mult_mat33_vec3(R_NF.Ptr(),
                                        tmp_vec3,
                                        nbetak.wPtr()
                                        );                                
        }
        else{                
                // This is the case when the given body is the root 
                // node and parent is ground
                const Transform& T_PM = Nd->data_->getTransformPM();
                
                FastOps::mult_mat33_vec3(T_PM.R().Ptr(),
                                        T_MF.P().Ptr(),
                                        tmp_vec3
                                        );
                
                double *betaptr = nbetak.wPtr();
                const double *ptr = T_PM.P().Ptr();
                
                betaptr[0] = tmp_vec3[0] + ptr[0];
                betaptr[1] = tmp_vec3[1] + ptr[1];
                betaptr[2] = tmp_vec3[2] + ptr[2];
        }
        
        int len = Nd->Child_.getLength();
        
        // Repeat procedure for all the bodies in the system
        for(int i=0;i<len;++i)
                recursiveUpdSystemNBetak(state,Nd->Child_[i]);
        
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::recursiveUpdSystemRotationMatrixNF(const State& state, RigidBodyNode *Nd){
        int len = Nd->Child_.getLength();
        for(int i=0;i<len;++i){
                RigidBodyNode *tmp = Nd->Child_.getNode(i);

                // prF_C_F = Parent_FM * Body_MF
                FastOps::mult_mat33_mat33(Nd->data_->getBodyTransformFM()[i].R().Ptr(),
                                        tmp->data_->getJointRotationMF().Ptr(),
                                        tmp->data_->wgetRotationFCF().wPtr()
                                        );

                // NF_parent*prF_C_F = NF_child
                FastOps::mult_mat33_mat33(Nd->data_->getRotationNF().Ptr(),
                                        tmp->data_->getRotationFCF().Ptr(),
                                        tmp->data_->wgetRotationNF().wPtr()
                                        );

                recursiveUpdSystemRotationMatrixNF(state, tmp);
        }
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::updStateFromPrescribedMotion(State& state){
        state.updStateFromPrescribedMotion();
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::updU(const BodyIndex& BI, const double *u, State& state)const{
        #ifdef RSIM_SAFE
                assert(BI.id() < nNodes_);
                assert(BI.id()>=0);
        #endif

        state.updU(u,BI.id());

}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::updQ(const BodyIndex& BI,const double *q,State& state)const{
        #ifdef RSIM_SAFE	
                assert(BI.id() < nNodes_);
                assert(BI.id()>=0);
        #endif
        state.updQ(q,BI.id());
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::updQ(const NodeIndex& NI,const double *q,State& state)const{
        #ifdef RSIM_SAFE        
                assert(NI.id() < nNodes_);
                assert(NI.id()>=0);
        #endif
        state.updQ(q,NI.id());
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::updOneQ(const BodyIndex& BI, const QIndex& qindex, const double& q, State& state)const{
        #ifdef RSIM_SAFE	
                Utils::rsim_assert(BI.id() < nNodes_,
                                "MultibodySystem::updOneQ(), BodyIndex is too high\n");
                Utils::rsim_assert(BI.id()>=0,
                                "MultibodySystem::updOneQ(), BodyIndex must be greater than or equal to zero\n");
        #endif
        state.updOneQ(q, BI.id(), qindex);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::updOneU(const BodyIndex& BI, const UIndex& uindex, const double& u, State& state)const{
        #ifdef RSIM_SAFE	
                assert(BI.id() < nNodes_);
                assert(BI.id()>=0);
        #endif
        state.updOneU(u, BI.id(), uindex);
}

////////////////////////////////////////////////////////////////////////

int MultibodySystem::getNBodies() const {
        return (nNodes_+1);
}

////////////////////////////////////////////////////////////////////////

int MultibodySystem::getNNodes() const {
        return nNodes_;
}

////////////////////////////////////////////////////////////////////////

const RigidBodyNode* MultibodySystem::getRoot() const {
        return Root_;
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::addPrescribedMotion(const BodyIndex& BI, PrescribedMotion* Ptr){
        if(BI.id()>=this->nNodes_){
                cout<<"\nMultibodySystem::addPrescribedMotion(const BodyIndex& BI, PrescribedMotion* Ptr): BodyIndex too high\n";
                exit(0);
        }
        BodyArray_[BI.id()]->data_->addPrescribedMotion(Ptr);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::stopPrescribedMotion(const BodyIndex& BI){
        BodyArray_[BI.id()]->data_->stopPrescribedMotion();
}

////////////////////////////////////////////////////////////////////////
                
void MultibodySystem::continuePrescribedMotion(const BodyIndex& BI){
        BodyArray_[BI.id()]->data_->continuePrescribedMotion();
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::addExternalSpatialForce(const BodyIndex& BI, ExternalSpatialForce *F){
        #ifdef RSIM_SAFE
        if(BI.id()>=this->nNodes_){
                cout<<"\nMultibodySystem::addExternalSpatialForce(const BodyIndex& BI, ExternalSpatialForce *F): BodyIndex too high\n";
                exit(0);
        }
        #endif
        BodyArray_[BI.id()]->data_->addExternalSpatialForce(F);
}

////////////////////////////////////////////////////////////////////////

MultibodySystem::MultibodySystem():	Root_(NULL),nNodes_(0),
                                        isTopologyMapSet_(false),
                                        SysAccl_(NULL){
        Stage::CurrentStage = Stage::EMPTY;
        set_new_handler(UnableToAllocateMemory);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::initInverseDynamics(){
        #ifdef RSIM_SAFE
        for(int i=0;i<nNodes_;++i){
                // Ensure that motion is prescribed on each joint.
                if(!BodyArray_[i]->data_->getJoint()->getIsMotionPrescribedAndActive()){
                        cout<<"MultibodySystem::initInverseDynamics(): Motion must be prescribed and active on each joint "
                        <<"to perform inverse dynamics. Use function MultibodySystem::addPrescribedMotion() or "<<
                        "MultibodySystem::continuePrescribedMotion(const BodyIndex& BI) \n";
                        exit(0);
                }
        }
        #endif
        
        this->initEOM();	
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::recursiveSweepForwardRNEA(const int& ChildID, State& state, RigidBodyNode *Nd){
        Nd->sweepForwardRNEA(this->SysAccl_,ChildID,state);
        const int& len = Nd->Child_.getLength();
        for(int i=0;i<len;++i){
                RigidBodyNode *node = Nd->Child_.getNode(i);
                this->recursiveSweepForwardRNEA(i,state,node);
        }
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::initEOM(){
        this->recursiveInitEOM(Root_);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::initABA(){
        this->recursiveInitForwardDynamics(Root_);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::recursiveInitForwardDynamics(RigidBodyNode *Nd){
        const int& len = Nd->Child_.getLength();
        Nd->data_->initForwardDynamics(len);
        for(int i=0;i<len;++i){
                RigidBodyNode *node = Nd->Child_.getNode(i);
                this->recursiveInitForwardDynamics(node);
        }
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::recursiveInitEOM(RigidBodyNode *Nd){
        Nd->data_->initEOM();
        int len = Nd->Child_.getLength();
        for(int i=0;i<len;++i){
                RigidBodyNode *node = Nd->Child_.getNode(i);
                this->recursiveInitEOM(node);
        }
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::sweepBackwardsABA(State& state){
        this->recursiveSweepBackwardsABA(0,Root_, state);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::recursiveSweepBackwardsABA(const int& ChildID,RigidBodyNode *Nd,State& state){
        Nd->data_->formEOM(this->SysAccl_);
        int len = Nd->Child_.getLength();
        for(int i=0;i<len;++i){
                RigidBodyNode *tmp = Nd->Child_.getNode(i);
                this->recursiveSweepBackwardsABA(i,tmp,state);
        }
        Nd->sweepBackwardsABA(ChildID,state);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::sweepBackwardRNEA(State& state){
        this->recursiveSweepBackwardRNEA(0,Root_,state);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::recursiveSweepBackwardRNEA(const int& ChildID,
                                                 RigidBodyNode *Nd, 
                                                 State& state){
        Nd->sweepForwardRNEA(this->SysAccl_,ChildID,state);
        int len = Nd->Child_.getLength();
        for(int i=0;i<len;++i){
                RigidBodyNode *tmp = Nd->Child_.getNode(i);
                this->recursiveSweepBackwardRNEA(i,tmp,state);
        }
        Nd->sweepBackwardRNEA(ChildID,state);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::sweepForwardABA(State& state){
        this->recursiveSweepForwardABA(0,Root_,state);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::recursiveSweepForwardABA(const int& ChildID,RigidBodyNode *Nd, State& state){
        Nd->data_->sweepForwardABA(Nd->Parent_, ChildID, state);
        int len = Nd->Child_.getLength();
        for(int i=0;i<len;++i)
                this->recursiveSweepForwardABA(i,Nd->Child_.getNode(i),state);
}

////////////////////////////////////////////////////////////////////////

MultibodySystem::~MultibodySystem(){
        if(nNodes_){
                recursiveDeleteNode(Root_);
        }
        delete SysAccl_;
        BodyArray_.clear();
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::recursiveDeleteNode(RigidBodyNode *Nd){
        delete Nd->data_;
        int len = Nd->Child_.getLength();
        for(int i=0;i<len;++i)
                recursiveDeleteNode(Nd->Child_.getNode(i));
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::addSystemAcceleration(const Vect3& GravityLikeAcceleration){
        this->SysAccl_ = new SystemAcceleration(GravityLikeAcceleration);
}

////////////////////////////////////////////////////////////////////////

void MultibodySystem::getConstraintForceActingOnBody(const BodyIndex& BI, double *Fc)const{
        const int& id = BI.id();
        #ifdef RSIM_SAFE
        if (id >= this->nNodes_){
                cout<<"\nvoid MultibodySystem::getConstraintForceActingOnBody(const BodyIndex& BI, double *Fc)const;\
                BodyIndex too high. Exiting..\n";
                exit(0);
        }
        #endif
        const RigidBodyNode *Nd = BodyArray_[id];
        Nd->data_->getConstraintForce(Fc);
}